# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import sys
import os
from logging import getLogger
import scipy
import scipy.linalg
import torch
from torch.autograd import Variable
from torch.nn import functional as F

from .utils import get_optimizer, load_embeddings, normalize_embeddings, export_embeddings
from .utils import clip_parameters
from .dico_builder import build_dictionary
from .evaluation.word_translation import DIC_EVAL_PATH, load_identical_char_dico, load_dictionary

import numpy as np
from nltk.stem.porter import PorterStemmer
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from multiprocessing.pool import Pool
from functools import partial
import time
logger = getLogger()
pt=PorterStemmer()


class Trainer(object):

    def __init__(self, src_emb, tgt_emb, mapping, discriminator, params):
        """
        Initialize trainer script.
        始终没弄明白src_dico和tgt_dico中存储的是什么
        其实是从各自的embedding中读取构建的字典，序号对应词，至于序号则是按读取的顺序编排
        原始的params中并没有字典，是在bulid_model方法中构建的
        """
        self.src_emb = src_emb
        self.tgt_emb = tgt_emb
        self.src_dico = params.src_dico
        self.tgt_dico = getattr(params, 'tgt_dico', None)
        self.mapping = mapping
        self.discriminator = discriminator
        self.params = params

        # optimizers
        if hasattr(params, 'map_optimizer'):
            optim_fn, optim_params = get_optimizer(params.map_optimizer)
            self.map_optimizer = optim_fn(mapping.parameters(), **optim_params)
        if hasattr(params, 'dis_optimizer'):
            optim_fn, optim_params = get_optimizer(params.dis_optimizer)
            self.dis_optimizer = optim_fn(discriminator.parameters(), **optim_params)
        else:
            assert discriminator is None

        # best validation score
        self.best_valid_metric = -1e12

        self.decrease_lr = False

    def get_dis_xy(self, volatile):
        """
        Get discriminator input batch / output target.
        """
        # select random word IDs
        bs = self.params.batch_size
        mf = self.params.dis_most_frequent
        assert mf <= min(len(self.src_dico), len(self.tgt_dico))
        src_ids = torch.LongTensor(bs).random_(len(self.src_dico) if mf == 0 else mf)
        tgt_ids = torch.LongTensor(bs).random_(len(self.tgt_dico) if mf == 0 else mf)
        if self.params.cuda:
            src_ids = src_ids.cuda()
            tgt_ids = tgt_ids.cuda()

        # get word embeddings
        src_emb = self.src_emb(Variable(src_ids, volatile=True))
        tgt_emb = self.tgt_emb(Variable(tgt_ids, volatile=True))
        src_emb = self.mapping(Variable(src_emb.data, volatile=volatile))
        tgt_emb = Variable(tgt_emb.data, volatile=volatile)

        # input / target
        x = torch.cat([src_emb, tgt_emb], 0)
        y = torch.FloatTensor(2 * bs).zero_()
        y[:bs] = 1 - self.params.dis_smooth
        y[bs:] = self.params.dis_smooth
        y = Variable(y.cuda() if self.params.cuda else y)

        return x, y

    def dis_step(self, stats):
        """
        Train the discriminator.
        """
        self.discriminator.train()

        # loss
        x, y = self.get_dis_xy(volatile=True)
        preds = self.discriminator(Variable(x.data))
        loss = F.binary_cross_entropy(preds, y)
        stats['DIS_COSTS'].append(loss.data[0])

        # check NaN
        if (loss != loss).data.any():
            logger.error("NaN detected (discriminator)")
            exit()

        # optim
        self.dis_optimizer.zero_grad()
        loss.backward()
        self.dis_optimizer.step()
        clip_parameters(self.discriminator, self.params.dis_clip_weights)

    def mapping_step(self, stats):
        """
        Fooling discriminator training step.
        """
        if self.params.dis_lambda == 0:
            return 0

        self.discriminator.eval()

        # loss
        x, y = self.get_dis_xy(volatile=False)
        preds = self.discriminator(x)
        loss = F.binary_cross_entropy(preds, 1 - y)
        loss = self.params.dis_lambda * loss

        # check NaN
        if (loss != loss).data.any():
            logger.error("NaN detected (fool discriminator)")
            exit()

        # optim
        self.map_optimizer.zero_grad()
        loss.backward()
        self.map_optimizer.step()
        self.orthogonalize()

        return 2 * self.params.batch_size

    def load_training_dico(self, dico_train):
        """
        Load training dictionary.
        字典里面存了三个东西，word2id,id2word,lang,同时前面2个都是字典
        """
        word2id1 = self.src_dico.word2id
        word2id2 = self.tgt_dico.word2id
        # print(word2id1)
        # sys.exit(0)

        # identical character strings
        if dico_train == "identical_char":
            self.dico = load_identical_char_dico(word2id1, word2id2)
        # use one of the provided dictionary
        elif dico_train == "default":
            filename = '%s-%s.txt' % (self.params.src_lang, self.params.tgt_lang)
            self.dico = load_dictionary(
                # os.path.join('../', filename),
                os.path.join(DIC_EVAL_PATH, filename),
                word2id1, word2id2
            )
        # dictionary provided by the user
        else:
            self.dico = load_dictionary(dico_train, word2id1, word2id2)

        # cuda
        if self.params.cuda:
            self.dico = self.dico.cuda()

    def build_dictionary(self):
        """
        Build a dictionary from aligned embeddings.
        下面的对词向量的转换，算归一化吧。
        具体的在bulid_dictionary中向量点乘就相当于二者之间的相似度
        """
        src_emb = self.mapping(self.src_emb.weight).data
        tgt_emb = self.tgt_emb.weight.data
        src_emb = src_emb / src_emb.norm(2, 1, keepdim=True).expand_as(src_emb)
        tgt_emb = tgt_emb / tgt_emb.norm(2, 1, keepdim=True).expand_as(tgt_emb)
        self.dico = build_dictionary(src_emb, tgt_emb, self.params)

    #添加的一个函数用于加载平行句子向量
    def get_senVec(self,infile):
        with open(infile,'r',encoding='utf-8') as r:
            sentence_vec=[]
            lines=r.readlines()
            for i,line in enumerate(lines):
                if i==0:
                    continue
                list_sen=line.strip('\n').strip().split()
                array_sen=np.array(list_sen).astype('float64')
                sentence_vec.append(array_sen)
        sen_all=np.array(sentence_vec)
        sen=torch.from_numpy(sen_all)
        return sen

    def tokenizer(self,lines,lang):
        stop_words=stopwords.words(lang)
        sentences=[]
        for line in lines:
            word_tokens=word_tokenize(line.strip('/n').strip())
            filtered_sentence=[w for w in word_tokens if w not in stop_words]
            porter_sentence=[pt.stem(v) for v in filtered_sentence]
            sentences.append(porter_sentence)
        return sentences

    def add_sentenceVec(self,infile,lang,embeddings):
        if lang=='english':
            word2id = self.src_dico.word2id
        else:
            word2id=self.tgt_dico.word2id
        with open(infile,'r',encoding='utf-8') as r:
            lines=r.readlines()
            sentences=self.tokenizer(lines,lang)
            sen_vecList=[]
            for i,line in enumerate(lines):
                print(i)
                vec=np.zeros(300)
                for word in line:
                    if word in word2id.keys():
                        id=word2id[word]
                    else:
                        continue
                    vec+=embeddings[id]
                sen_vecList.append(vec)
        return torch.from_numpy(np.array(sen_vecList))

    def procrustes(self):
        """
        Find the best orthogonal matrix mapping using the Orthogonal Procrustes problem
        https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem
        torch.mm和torch.matmul结果一致，但是前者不能广播，后面可以广播。不太理解。
        A和B分别是训练的词对应的矩阵，已经从原本的序号还原成词向量了,A和B的形状都是[训练词的数目，词向量维度]
        mapping = nn.Linear(params.emb_dim, params.emb_dim, bias=False)
        mapping构建的一个线性连接层，即y=wx+b,但是这里bias设置为false，故mapping是一个线性变换，y=wx
        svd分解之后得出的都是numpy矩阵
        U的形状是[300,300],S的形状是[300,],V_t的形状是[300,300],分别是二维，一维，二维
        """
        # if n>0:
        #     A = self.src_emb.weight.data[self.dico[:, 0]]
        #     # print(A.size())
        #     # sys.exit(0)
        #     B = self.tgt_emb.weight.data[self.dico[:, 1]]
        # else:
        #     A=self.get_senVec('E:/my experiment/first_task/en_sentence.txt')
        #     B=self.get_senVec('E:/my experiment/first_task/it_sentence.txt')
        # A=self.add_sentenceVec('/media/liyulong/study/my experiment/first_task/en_sentence.txt','english',self.src_emb.weight.data)
        # B=self.add_sentenceVec('/media/liyulong/study/my experiment/first_task/it_sentence.txt','italian',self.tgt_emb.weight.data)
        # A = A / A.norm(2, 1, keepdim=True).expand_as(A)
        # B = B / B.norm(2, 1, keepdim=True).expand_as(B)
        A = self.src_emb.weight.data[self.dico[:, 0]]
        B = self.tgt_emb.weight.data[self.dico[:, 1]]
        W = self.mapping.weight.data
        print(W)
        M = B.transpose(0, 1).mm(A).cpu().numpy()
        U, S, V_t = scipy.linalg.svd(M, full_matrices=True)
        # print(U.shape)
        # print(S)
        # print(V_t.shape)
        # sys.exit(0)
        W.copy_(torch.from_numpy(U.dot(V_t)).type_as(W))


    def orthogonalize(self):
        """
        Orthogonalize the mapping.
        """
        if self.params.map_beta > 0:
            W = self.mapping.weight.data
            beta = self.params.map_beta
            W.copy_((1 + beta) * W - beta * W.mm(W.transpose(0, 1).mm(W)))

    def update_lr(self, to_log, metric):
        """
        Update learning rate when using SGD.
        """
        if 'sgd' not in self.params.map_optimizer:
            return
        old_lr = self.map_optimizer.param_groups[0]['lr']
        new_lr = max(self.params.min_lr, old_lr * self.params.lr_decay)
        if new_lr < old_lr:
            logger.info("Decreasing learning rate: %.8f -> %.8f" % (old_lr, new_lr))
            self.map_optimizer.param_groups[0]['lr'] = new_lr

        if self.params.lr_shrink < 1 and to_log[metric] >= -1e7:
            if to_log[metric] < self.best_valid_metric:
                logger.info("Validation metric is smaller than the best: %.5f vs %.5f"
                            % (to_log[metric], self.best_valid_metric))
                # decrease the learning rate, only if this is the
                # second time the validation metric decreases
                if self.decrease_lr:
                    old_lr = self.map_optimizer.param_groups[0]['lr']
                    self.map_optimizer.param_groups[0]['lr'] *= self.params.lr_shrink
                    logger.info("Shrinking the learning rate: %.5f -> %.5f"
                                % (old_lr, self.map_optimizer.param_groups[0]['lr']))
                self.decrease_lr = True

    def save_best(self, to_log, metric):
        """
        Save the best model for the given validation metric.
        """
        # best mapping for the given validation criterion
        if to_log[metric] > self.best_valid_metric:
            # new best mapping
            self.best_valid_metric = to_log[metric]
            logger.info('* Best value for "%s": %.5f' % (metric, to_log[metric]))
            # save the mapping
            W = self.mapping.weight.data.cpu().numpy()
            path = os.path.join(self.params.exp_path, 'best_mapping.pth')
            logger.info('* Saving the mapping to %s ...' % path)
            torch.save(W, path)

    def reload_best(self):
        """
        Reload the best mapping.
        """
        path = os.path.join(self.params.exp_path, 'best_mapping.pth')
        logger.info('* Reloading the best model from %s ...' % path)
        # reload the model
        assert os.path.isfile(path)
        to_reload = torch.from_numpy(torch.load(path))
        W = self.mapping.weight.data
        assert to_reload.size() == W.size()
        W.copy_(to_reload.type_as(W))

    def export(self):
        """
        Export embeddings.
        """
        params = self.params

        # load all embeddings
        logger.info("Reloading all embeddings for mapping ...")
        params.src_dico, src_emb = load_embeddings(params, source=True, full_vocab=True)
        params.tgt_dico, tgt_emb = load_embeddings(params, source=False, full_vocab=True)

        # apply same normalization as during training
        normalize_embeddings(src_emb, params.normalize_embeddings, mean=params.src_mean)
        normalize_embeddings(tgt_emb, params.normalize_embeddings, mean=params.tgt_mean)

        # map source embeddings to the target space
        bs = 4096
        logger.info("Map source embeddings to the target space ...")
        for i, k in enumerate(range(0, len(src_emb), bs)):
            x = Variable(src_emb[k:k + bs], volatile=True)
            src_emb[k:k + bs] = self.mapping(x.cuda() if params.cuda else x).data.cpu()

        # write embeddings to the disk
        export_embeddings(src_emb, tgt_emb, params)
